{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 分布式训练\n",
    "## 概述\n",
    "[tf.distribute.Strategy](https://tensorflow.google.cn/api_docs/python/tf/distribute/Strategy) 是一个用于跨多个 GPU、多台机器或 TPU 分布训练的 tensorflow API。使用这个 API，你可以用最少的代码更改分发现有的模型和训练代码。\n",
    "\n",
    "tf.distribute.Strategy 的设计考虑了以下关键目标：\n",
    "- 易于使用和支持多个用户层，包括研究人员、ML 工程师等。\n",
    "- 提供良好的开箱即用性能。\n",
    "- 策略之间容易进行切换。\n",
    "\n",
    "tf.distribute.Strategy 可以与类似 Keras 的高级 API 一起使用，还可以用于分发自定义训练循环（通常，还可以使用 tensorflow 进行任何计算）。\n",
    "\n",
    "在 tensorflow2.0 中，可以在 eager execution 下执行你的程序，也可以使用 tf.function 在图形中执行程序。tf.distribute.Strategy 支持这两种执行模式。尽管我们在本指南中大部分时间都在讨论训练，但此 API 也可以用于在不同平台上分发评估和预测。\n",
    "\n",
    "你可以在很少修改代码的情况下使用 tf.distribute.Strategy，因为 tensorflow 已经更改了它的底层组件，使其能够感知策略。这包括变量、层、模型、优化器、度量、摘要和检查点。\n",
    "\n",
    "在本指南中，我们将解释各种类型的策略以及如何在不同的情况下使用它们。\n",
    "\n",
    "> 注意：为了深层次地理解这个概念，你可以阅读 [this deep-dive presentation](https://youtu.be/jKV53r9-H14)。如果你打算编写自己的训练循环，强烈推荐这一点。\n",
    "\n",
    "## 设置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 策略的类型\n",
    "tf.distribute.Strategy 打算根据不同的轴覆盖多种用例。目前支持其中一些组合，将来还将添加其他组合。其中一些轴是：\n",
    "- 同步与异步训练（Synchronous vs asynchronous training）：这是两种常见的数据并行分布式训练方法。在同步训练中，所有的工作者在不同的输入数据片段上进行同步训练，并在每一步聚集梯度。在异步训练中，所有的工作者都独立地对输入数据进行训练，并异步地更新变量。通常，通过参数服务器（参数服务器是个编程框架，用于方便分布式并行程序的编写，其中重点是对大规模参数的分布式存储和协同的支持。）体系结构，所有 reduce 和 async 都支持同步培训。\n",
    "- 硬件平台（Hardware platform）：你可能希望将培训扩展到一台计算机上的多个GPU，或网络中的多台计算机（每个 GPU 为 0 或多个）或云 TPU 上。\n",
    "\n",
    "为了支持这些用例，有六种策略可用。在下一节中，我们将解释目前 TF2 中支持哪些场景。下面是一个快速概述：\n",
    "\n",
    "| 训练 API | 镜像策略 | TPU 策略 | 多工作者镜像策略 | 中央存储策略 | 参数服务器策略 | 单体策略 |\n",
    "| - | - | - | - | - | - | - |\n",
    "| Keras API | 支持 | 实验性支持 | 实验性支持 | 实验性支持 | 计划 2.0 后支持 | 支持 |\n",
    "| 自定义训练循环 | 实验性支持 | 实验性支持 | 计划 2.0 后支持 | 目前不支持 | 支持 | 支持 |\n",
    "| Estimator API | 限制性支持 | 不支持 | 限制性支持 | 限制性支持 | 限制性支持 | 限制性支持 |\n",
    "\n",
    "> 注意：对 Estimator 支持受到限制。基本训练和评估都是实验性的，没有实现脚手架等先进功能。我们建议在不涉及用例的情况下使用 Keras 或自定义训练循环。\n",
    "\n",
    "### 镜像策略\n",
    "[tf.distribute.MirroredStrategy](https://tensorflow.google.cn/api_docs/python/tf/distribute/MirroredStrategy) 支持在一台机器上对多个 GPU 进行同步分布式训练。它为每个 GPU 设备创建一个副本。模型中的每个变量都在所有副本中镜像。这些变量一起构成一个称为 MirroredVariable 的概念变量。通过应用相同的更新，这些变量彼此保持同步。\n",
    "\n",
    "所有高效的 reduce 算法用于在设备之间传递变量更新。All-reduce 将所有涉笔上的聚合张量相加，使它们在每个设备上可用。它是一种高效的融合算法，可以显著降低同步开销。根据设备之间可用的通信类型，有许多可用的 reduce 算法和实现。默认情况下，它使用 NVIDIA NCCL 作为 All-reduce实现。你可以从 tensorflow 提供的其他选项中选择，也可以自己编写。\n",
    "\n",
    "以下是创建 MirroredStrategy 的最简单方法："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:There are non-GPU devices in `tf.distribute.Strategy`, not using nccl allreduce.\n",
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0',)\n"
     ]
    }
   ],
   "source": [
    "mirrored_strategy = tf.distribute.MirroredStrategy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这将创建一个 MirroredStrategy 实例，该实例将使用 tensorflow 可见的所有 GPU，并使用 NCCL 作为跨设备通信。\n",
    "\n",
    "如果你希望仅使用计算机上的部分 GPU，可以这样做："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mirrored_strategy = tf.distribute.MirroredStrategy(devices=[\"/gpu:0\", \"/gpu:1\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "如果希望重写跨设备通信，可以使用 cross-device-ops 参数通过提供的 [tf.distribute.CrossDeviceOps](https://tensorflow.google.cn/api_docs/python/tf/distribute/CrossDeviceOps) 实例来执行此操作。目前，[tf.distribute.HierarchicalCopyAllReduce](https://tensorflow.google.cn/api_docs/python/tf/distribute/HierarchicalCopyAllReduce) 和 [tf.distribute.ReductionToOneDevice](https://tensorflow.google.cn/api_docs/python/tf/distribute/ReductionToOneDevice) 是除默认选项 [tf.distribute.NcclAllReduce](https://tensorflow.google.cn/api_docs/python/tf/distribute/NcclAllReduce) 外的另外两个选项。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0',)\n"
     ]
    }
   ],
   "source": [
    "mirrored_strategy = tf.distribute.MirroredStrategy(\n",
    "    cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 中央存储策略\n",
    "[tf.distribute.experimental.CentralStorageStrategy](https://tensorflow.google.cn/api_docs/python/tf/distribute/experimental/CentralStorageStrategy) 也进行同步训练。变量不是镜像，而是放在 CPU 上，操作在所有本地 GPU 上复制。如果只有一个 GPU，则所有变量和操作都将放在该 GPU 上。\n",
    "\n",
    "通过以下方式创建 CentralStorageStrategy 实例："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:ParameterServerStrategy (CentralStorageStrategy if you are using a single machine) with compute_devices = ('/device:CPU:0',), variable_device = '/device:CPU:0'\n"
     ]
    }
   ],
   "source": [
    "central_storage_strategy = tf.distribute.experimental.CentralStorageStrategy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 多工作者镜像策略\n",
    "[tf.distribute.experimental.MultiWorkerMirroredStrategy](https://tensorflow.google.cn/api_docs/python/tf/distribute/experimental/MultiWorkerMirroredStrategy) 与 MirroredStrategy 非常相似。它实现了跨多个工作者的同步分布式培训，每个工作者可能有多个 GPU。与 MirroredStrategy 类似，它在所有工作进程中的每个设备上创建模型中所有变量的副本。\n",
    "\n",
    "它使用 [CollectiveOps](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/collective_ops.py) 作为多工作进程 all-reduce 通信方法，用于保持变量同步。集合运算是 tensotflow 图中的单个运算，它可以根据硬件、网络拓扑结构和张量大小在 tensorflow 运行时自动选择全归约算法。\n",
    "\n",
    "它还实现了额外的性能优化。例如，它包含一个静态优化，将小张量上的多重全归约转换为大张量上的较少全归约。此外，tensotflow 正在设计它的一个插件架构——所以在未来，你将能够插件算法，更好地为你的硬件进行调整。注意，集体操作还实现其他集体操作，如广播和所有聚集。\n",
    "\n",
    "以下是创建多工作 MultiWorkerMirroredStrategy 的最简单方法："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.\n",
      "INFO:tensorflow:Using MirroredStrategy with devices ('/device:CPU:0',)\n",
      "INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:CPU:0',), communication = CollectiveCommunication.AUTO\n"
     ]
    }
   ],
   "source": [
    "multiworker_strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "MultiWorkerMirroredStrategy 目前允许你在两种不同的集体操作实现之间进行选择。[CollectiveCommunication.RING](https://tensorflow.google.cn/api_docs/python/tf/distribute/experimental/CollectiveCommunication#RING) 使用 gRPC 作为通信层实现基于环的集合。[CollectiveCommunication.NCCL](https://tensorflow.google.cn/api_docs/python/tf/distribute/experimental/CollectiveCommunication#NCCL) 使用 [Nvidia 的 NCCL](https://developer.nvidia.com/nccl) 来实现集体操作。[CollectiveCommunication.AUTO]https://tensorflow.google.cn/api_docs/python/tf/distribute/experimental/CollectiveCommunication#AUTO() 将选择推迟到运行时。集体操作实现的最佳选择取决于 GPU 的数量和种类，以及集群中的网络互连。可以按以下方式指定它们："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.\n",
      "INFO:tensorflow:Using MirroredStrategy with devices ('/device:CPU:0',)\n",
      "WARNING:tensorflow:Enabled NCCL communication but no GPUs detected/specified.\n",
      "INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:CPU:0',), communication = CollectiveCommunication.NCCL\n"
     ]
    }
   ],
   "source": [
    "multiworker_strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(\n",
    "    tf.distribute.experimental.CollectiveCommunication.NCCL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "与多 GPU 培训相比，多工作者训练的关键区别之一是多工作者设置。TF_CONFIG 环境变量是 tensorflow 中为作为集群一部分的每个工作机指定集群配置的标准方法。了解有关设置 [setting up TF_CONFIG](https://tensorflow.google.cn/guide/distributed_training#TF_CONFIG) 更多信息。\n",
    "\n",
    "> 注意：这个策略是 [实验性](https://tensorflow.google.cn/guide/versions#what_is_not_covered) 的，因为目前正在改进它，并使其适用于更多的场景。作为这项工作的一部分，请期待 API 在未来发生变化。\n",
    "\n",
    "### TPU 策略\n",
    "[tf.distribute.experimental.TPUStrategy](https://tensorflow.google.cn/api_docs/python/tf/distribute/experimental/TPUStrategy) 允许你在张量处理单元（TPU）上进行张量训练。TPU 是 Google 的专用 ASIC，旨在显著加速机器学习工作负载。它们可以在 Google Colab、[TensorFlow Research Cloud](https://tensorflow.google.cn/tfrc) 和 [Cloud TPU](https://cloud.google.com/tpu) 上找到。\n",
    "\n",
    "在分布式训练体系结构方面，TPUStrategy 是相同的镜像策略，它实现了同步的分布式训练。TPU 在TPUStrategy 中使用的多个 TPU 核心上提供自己的高效 all-reduce 和其他集合操作的实现。\n",
    "\n",
    "下面是如何实例化 TPUStrategy的：\n",
    "\n",
    "> 注意：要在 Colab 中运行此代码，你应该选择 TPU 作为 Colab 运行时环境。很快就会有一个教程，演示如何使用 TPUStrategy。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'tpu_address' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-7-a0ff54b7a52b>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      1\u001b[0m cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(\n\u001b[1;32m----> 2\u001b[1;33m     tpu=tpu_address)\n\u001b[0m\u001b[0;32m      3\u001b[0m \u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mconfig\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mexperimental_connect_to_cluster\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mcluster_resolver\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      4\u001b[0m \u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mtpu\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mexperimental\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0minitialize_tpu_system\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mcluster_resolver\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      5\u001b[0m \u001b[0mtpu_strategy\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mdistribute\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mexperimental\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mTPUStrategy\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mcluster_resolver\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mNameError\u001b[0m: name 'tpu_address' is not defined"
     ]
    }
   ],
   "source": [
    "cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(\n",
    "    tpu=tpu_address)\n",
    "tf.config.experimental_connect_to_cluster(cluster_resolver)\n",
    "tf.tpu.experimental.initialize_tpu_system(cluster_resolver)\n",
    "tpu_strategy = tf.distribute.experimental.TPUStrategy(cluster_resolver)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TPUClusterResolver 实例能帮助定位 TPU。在 Colab 中，你不需要为它指定任何参数。\n",
    "\n",
    "如果要在 Cloud TPU 使用它：\n",
    "- 必须在参数 TPU 中指定 TPU 资源的名称。\n",
    "- 必须在程序开始时显式初始化 TPU 系统。在使用 TPU 进行计算之前，这是必需的。初始化 TPU 系统也会清除 TPU 内存，因此必须先完成此步骤，以避免丢失状态。\n",
    "\n",
    "> 注意：这个策略是实验性的，因为 tensotflow 目前正在改进它，并使其适用于更多的场景。作为这项工作的一部分，请期待 API 在未来发生的变化。\n",
    "\n",
    "### 参数服务器策略\n",
    "[tf.distribute.experimental.ParameterServerStrategy](https://tensorflow.google.cn/api_docs/python/tf/distribute/experimental/ParameterServerStrategy) 支持多台机器上的参数服务器培训。在这个设置中，有些机器被指定为工作机，有些被指定为参数服务器。模型的每个变量都放在一个参数服务器上。计算被复制到所有工作机的所有 GPU 上。\n",
    "\n",
    "在代码方面，它看起来与其他策略类似："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "Cluster spec must be non-empty in `cluster_resolver`.",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-8-3345f5629019>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[1;32m----> 1\u001b[1;33m \u001b[0mps_strategy\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mdistribute\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mexperimental\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mParameterServerStrategy\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;32mf:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\distribute\\parameter_server_strategy.py\u001b[0m in \u001b[0;36m__init__\u001b[1;34m(self, cluster_resolver)\u001b[0m\n\u001b[0;32m    111\u001b[0m       \u001b[0mcluster_resolver\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mTFConfigClusterResolver\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    112\u001b[0m     \u001b[1;32mif\u001b[0m \u001b[1;32mnot\u001b[0m \u001b[0mcluster_resolver\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcluster_spec\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 113\u001b[1;33m       \u001b[1;32mraise\u001b[0m \u001b[0mValueError\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m\"Cluster spec must be non-empty in `cluster_resolver`.\"\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    114\u001b[0m     extended = ParameterServerStrategyExtended(\n\u001b[0;32m    115\u001b[0m         self, cluster_resolver=cluster_resolver)\n",
      "\u001b[1;31mValueError\u001b[0m: Cluster spec must be non-empty in `cluster_resolver`."
     ]
    }
   ],
   "source": [
    "ps_strategy = tf.distribute.experimental.ParameterServerStrategy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "对于多工作机训练，TF_CONFIG 需要指定集群中参数服务器和过机的配置，你可以在 [TF_CONFIG](https://tensorflow.google.cn/guide/distributed_training#TF_CONFIG) 中阅读更多信息。\n",
    "\n",
    "### 单设备策略\n",
    "[tf.distribute.OneDeviceStrategy](https://tensorflow.google.cn/api_docs/python/tf/distribute/OneDeviceStrategy) 在单个设备上运行。此策略将在其作用域中创建的任何变量放置在指定设备上。通过此策略分发的输入将预取到指定设备。此外，通过 strategy.run 调用的任何函数也将放置在指定的设备上。\n",
    "\n",
    "在切换到实际分布到多个设备/机器的其他策略之前，可以使用此策略测试代码。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "strategy = tf.distribute.OneDeviceStrategy(device=\"/gpu:0\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "至此，我们已经讨论了可用的不同策略以及如何实例化它们。在接下来的几节中，我们将讨论使用它们分发训练的不同方式。我们将在本指南中显示简短的代码片段，并链接到完整的教程，你可以端到端地运行。\n",
    "\n",
    "## 在 Keras 使用 tf.distribute.Strategy\n",
    "tensorflow 已经将 tf.distribute.Strategy 集成到 tf.keras 中，这是 tensorflow 对 Keras API 规范的实现。tf.Keras 是一个高级 API，用于构建和训练模型。通过集成到 tf.keras 后端，tensorflow 能使你无缝地分发在 Keras 训练框架中编写的训练。\n",
    "\n",
    "以下是你需要在代码中更改的内容：\n",
    "- 创建适当的 tf.distribute.Strategy 实例\n",
    "- 将 Keras 模型的创建和编译移动到 strategy.scope中 。\n",
    "\n",
    "tensorflow 支持所有类型的 Keras models-sequential、函数和子类。\n",
    "\n",
    "下面是一段代码，用于一个非常简单的带有密集层的 Keras 模型："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:There are non-GPU devices in `tf.distribute.Strategy`, not using nccl allreduce.\n",
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0',)\n"
     ]
    }
   ],
   "source": [
    "mirrored_strategy = tf.distribute.MirroredStrategy()\n",
    "with mirrored_strategy.scope():\n",
    "  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])\n",
    "  model.compile(loss='mse', optimizer='sgd')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在本例中，我们使用了 MirroredStrategy，因此可以在具有多个 GPU 的计算机上运行它。.scope() 表示要运行的代码的哪些部分是分布式的。在这个范围内创建一个模型允许我们创建镜像变量而不是常规变量。在这个范围内编译可以让我们知道用户打算使用这个策略来训练这个模型。一旦设置好，你就可以像正常情况一样安装你的模型。MirroredStrategy 负责在可用的 GPU 上复制模型的训练，聚合梯度等等。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train for 10 steps\n",
      "Epoch 1/2\n",
      "10/10 [==============================] - 3s 340ms/step - loss: 0.2220\n",
      "Epoch 2/2\n",
      "10/10 [==============================] - 0s 3ms/step - loss: 0.0981\n",
      "10/10 [==============================] - 1s 107ms/step - loss: 0.0609\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0.06093804910778999"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(10)\n",
    "model.fit(dataset, epochs=2)\n",
    "model.evaluate(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这里我们使用 [tf.data.Dataset](https://tensorflow.google.cn/api_docs/python/tf/data/Dataset) 提供训练和评估输入。也可以使用 numpy 数组："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 100 samples\n",
      "Epoch 1/2\n",
      "100/100 [==============================] - 0s 3ms/sample - loss: 0.0434\n",
      "Epoch 2/2\n",
      "100/100 [==============================] - 0s 90us/sample - loss: 0.0192\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.keras.callbacks.History at 0x200c0817988>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "inputs, targets = np.ones((100, 1)), np.ones((100, 1))\n",
    "model.fit(inputs, targets, epochs=2, batch_size=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在这两种情况下（dataset 或 numpy），每批给定的输入在多个副本中平均分配。例如，如果对 2 个 GPU 使用 MirroredStrategy，那么大小为 10 的每个批将在 2 个 GPU 之间进行划分，每个 GPU 在每个步骤中接收 5 个输入示例。随着你增加更多的 GPU，每个 epoch 都会训练得更快。通常，你需要再添加更多加速器时增加批处理大小，以便有效地利用额外的计算能力。你还需要根据模型重新调整你的学习率。你可以使用 strategy.num_replications_in_sync 获取副本数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 使用副本的数量计算全局批大小\n",
    "BATCH_SIZE_PER_REPLICA = 5\n",
    "global_batch_size = (BATCH_SIZE_PER_REPLICA *\n",
    "                     mirrored_strategy.num_replicas_in_sync)\n",
    "dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100)\n",
    "dataset = dataset.batch(global_batch_size)\n",
    "\n",
    "LEARNING_RATES_BY_BATCH_SIZE = {5: 0.1, 10: 0.15}\n",
    "learning_rate = LEARNING_RATES_BY_BATCH_SIZE[global_batch_size]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 现在支持什么？\n",
    "在 TF2.0 版本中，Keras 支持 MirroredStrategy、TPUStrategy、CentralStorageStrategy 和 MultiWorkerMirroredStrategy。除了镜像策略，其他策略目前都处于试验阶段，随时可能改变。对其他策略的支持即将到来。API 和使用方法与上述完全相同。\n",
    "\n",
    "| 训练 API | MirroredStrategy | TPUStrategy | MultiWorkerMirroredStrategy | CentralStorageStrategy | ParameterServerStrategy | OneDeviceStrategy |\n",
    "| - | - | - | - | - |- | - |\n",
    "| Keras API | 支持 | 实验性支持 | 实验性支持 | 实验性支持 | 计划 2.0 后支持 | 支持 |\n",
    "\n",
    "### 示例和教程\n",
    "以下是一个教程和示例列表，说明了上述与Keras端到端的集成：\n",
    "1. 使用 MirroredStrategy 训练 [MNIST](https://tensorflow.google.cn/tutorials/distribute/keras) 的教程。\n",
    "2. 使用 MirroredStrategy 的 ImageNet 数据进行正式的 [ResNet50](https://github.com/tensorflow/models/blob/master/official/vision/image_classification/resnet_imagenet_main.py) 训练。\n",
    "3. [ResNet50](https://github.com/tensorflow/models/blob/master/official/vision/image_classification/resnet_imagenet_main.py) 使用 TPUStrategy 在 Cloud  CPU 上训练 Imagenet 数据。\n",
    "4. 使用 MultiWorkerMirroredStrategy 训练 MNIST 的 [教程](https://tensorflow.google.cn/tutorials/distribute/multi_worker_with_keras)。\n",
    "5. [NCF](https://github.com/tensorflow/models/blob/master/official/recommendation/ncf_keras_main.py) 使用 MirroredStrategy 进行训练。\n",
    "6. 使用镜像策略训练 [Transformer](https://github.com/tensorflow/models/blob/master/official/nlp/transformer/transformer_main.py)。\n",
    "\n",
    "## 在自定义训练循环中使用 tf.distribute.Strategy\n",
    "如你所见，将 tf.distribute.Strategy 与高级 API（Estimator和Keras）结合使用只需要更改几行代码。再花点功夫，你也可以使用 tf.distribute.Strategy 和定制的训练循环。\n",
    "\n",
    "如果你需要比 Estimator 或 Keras 更大的灵活性和对训练循环的控制，可以编写自定义的训练循环。例如，当使用 GAN 时，你可能希望每轮执行不同数量的生成器或鉴别器步骤。同样，高层次的框架也不太适合强化学习训练。\n",
    "\n",
    "为了支持定制的训练循环，我们通过 tf.distribute.Strategy 类提供了一组核心方法。使用这些可能需要在开始时对代码进行少量的重组，但是一旦完成，你应该能够通过更改策略实例在 GPU、TPU 和多台计算机之间进行切换。\n",
    "\n",
    "在这里，我们将展示一个简短的代码段来说明这个用例，这个用例是一个简单的训练示例，使用的是与前面相同的 Keras 模型。\n",
    "\n",
    "首先，我们在策略的范围内创建模型和优化器。这可以确保使用模型和优化器创建的任何变量都是镜像变量。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "with mirrored_strategy.scope():\n",
    "  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])\n",
    "  optimizer = tf.keras.optimizers.SGD()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接下来，我们创建输入数据集，并调用 [tf.distribute.Strategy.experimental_distribute_dataset](https://tensorflow.google.cn/api_docs/python/tf/distribute/Strategy#experimental_distribute_dataset) 来基于该策略分发数据集。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(1000).batch(\n",
    "    global_batch_size)\n",
    "dist_dataset = mirrored_strategy.experimental_distribute_dataset(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "然后，我们定义训练的一个步骤。我们将使用 [tf.GradientTape](https://tensorflow.google.cn/api_docs/python/tf/GradientTape) 计算梯度，并使用优化器应用这些梯度来更新模型的变量。要分发此培训步骤，我们将一个函数 step_fn 放入，并将其传递给tf.distbute.Strategy.run 以及我们从之前创建的 dist_dataset 获得的数据集输入："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def train_step(dist_inputs):\n",
    "  def step_fn(inputs):\n",
    "    features, labels = inputs\n",
    "\n",
    "    with tf.GradientTape() as tape:\n",
    "      # training=True 只在训练与推理(例如，退出)之间存在具有不同行为的层时才需要。\n",
    "      logits = model(features, training=True)\n",
    "      cross_entropy = tf.nn.softmax_cross_entropy_with_logits(\n",
    "          logits=logits, labels=labels)\n",
    "      loss = tf.reduce_sum(cross_entropy) * (1.0 / global_batch_size)\n",
    "\n",
    "    grads = tape.gradient(loss, model.trainable_variables)\n",
    "    optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))\n",
    "    return cross_entropy\n",
    "\n",
    "  per_example_losses = mirrored_strategy.experimental_run_v2(step_fn, args=(dist_inputs,))\n",
    "  mean_loss = mirrored_strategy.reduce(\n",
    "      tf.distribute.ReduceOp.MEAN, per_example_losses, axis=0)\n",
    "  return mean_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在上面的代码中还需要注意以下几点：\n",
    "1. 我们使用 [tf.nn.softmax_cross_entropy_and_logits](https://tensorflow.google.cn/api_docs/python/tf/nn/softmax_cross_entropy_with_logits) 计算损失。然后我们将总损失按全局批量大小进行缩放。这一点很重要，因为所有副本都是同步训练的，并且每个训练步骤中的示例数都是全局批处理。因此，损失需要除以全局批大小，而不是副本（本地）批大小。\n",
    "2. 我们使用 [tf.distribute.Strategy.reduce](https://tensorflow.google.cn/api_docs/python/tf/distribute/Strategy#reduce) API 聚合 tf.distribute.Strategy.run 返回的结果。tf.distribute.Strategy.run 从策略中的每个本地副本返回结果，有多种方法可以使用此结果。你可以减少它们以获得聚合值。你还可以执行 [tf.distribute.Strategy.experimental_local_results](https://tensorflow.google.cn/api_docs/python/tf/distribute/Strategy#experimental_local_results) 以获取结果中包含的值列表，每个本地副本一个。\n",
    "3. 当在分布策略范围内调用 apply_gradients 时，将修改其行为。具体来说，在同步训练期间对每个并行实例应用梯度之前，它对梯度的所有副本执行求和。\n",
    "\n",
    "最后，定义了训练步骤后，我们可以在 dist_dataset 上迭代并循环运行训练："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "in converted code:\n\n    <ipython-input-16-f1680ec5eaf9>:17 train_step  *\n        per_example_losses = mirrored_strategy.run(step_fn, args=(dist_inputs,))\n\n    AttributeError: 'MirroredStrategy' object has no attribute 'run'\n",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-29-3472ab066f79>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[1;32mwith\u001b[0m \u001b[0mmirrored_strategy\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mscope\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      2\u001b[0m   \u001b[1;32mfor\u001b[0m \u001b[0minputs\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mdist_dataset\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 3\u001b[1;33m     \u001b[0mprint\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mtrain_step\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;32mf:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\eager\\def_function.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, *args, **kwds)\u001b[0m\n\u001b[0;32m    566\u001b[0m         \u001b[0mxla_context\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mExit\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    567\u001b[0m     \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 568\u001b[1;33m       \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_call\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    569\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    570\u001b[0m     \u001b[1;32mif\u001b[0m \u001b[0mtracing_count\u001b[0m \u001b[1;33m==\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_get_tracing_count\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mf:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\eager\\def_function.py\u001b[0m in \u001b[0;36m_call\u001b[1;34m(self, *args, **kwds)\u001b[0m\n\u001b[0;32m    613\u001b[0m       \u001b[1;31m# This is the first call of __call__, so we have to initialize.\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    614\u001b[0m       \u001b[0minitializers\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m[\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 615\u001b[1;33m       \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_initialize\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mkwds\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0madd_initializers_to\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0minitializers\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    616\u001b[0m     \u001b[1;32mfinally\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    617\u001b[0m       \u001b[1;31m# At this point we know that the initialization is complete (or less\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mf:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\eager\\def_function.py\u001b[0m in \u001b[0;36m_initialize\u001b[1;34m(self, args, kwds, add_initializers_to)\u001b[0m\n\u001b[0;32m    495\u001b[0m     self._concrete_stateful_fn = (\n\u001b[0;32m    496\u001b[0m         self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access\n\u001b[1;32m--> 497\u001b[1;33m             *args, **kwds))\n\u001b[0m\u001b[0;32m    498\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    499\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0minvalid_creator_scope\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0munused_args\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0munused_kwds\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mf:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\eager\\function.py\u001b[0m in \u001b[0;36m_get_concrete_function_internal_garbage_collected\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   2387\u001b[0m       \u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mkwargs\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;32mNone\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;32mNone\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2388\u001b[0m     \u001b[1;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_lock\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 2389\u001b[1;33m       \u001b[0mgraph_function\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0m_\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0m_\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_maybe_define_function\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   2390\u001b[0m     \u001b[1;32mreturn\u001b[0m \u001b[0mgraph_function\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2391\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mf:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\eager\\function.py\u001b[0m in \u001b[0;36m_maybe_define_function\u001b[1;34m(self, args, kwargs)\u001b[0m\n\u001b[0;32m   2701\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2702\u001b[0m       \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_function_cache\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmissed\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0madd\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mcall_context_key\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 2703\u001b[1;33m       \u001b[0mgraph_function\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_create_graph_function\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   2704\u001b[0m       \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_function_cache\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mprimary\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mcache_key\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mgraph_function\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2705\u001b[0m       \u001b[1;32mreturn\u001b[0m \u001b[0mgraph_function\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mf:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\eager\\function.py\u001b[0m in \u001b[0;36m_create_graph_function\u001b[1;34m(self, args, kwargs, override_flat_arg_shapes)\u001b[0m\n\u001b[0;32m   2591\u001b[0m             \u001b[0marg_names\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0marg_names\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2592\u001b[0m             \u001b[0moverride_flat_arg_shapes\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0moverride_flat_arg_shapes\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 2593\u001b[1;33m             capture_by_value=self._capture_by_value),\n\u001b[0m\u001b[0;32m   2594\u001b[0m         \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_function_attributes\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2595\u001b[0m         \u001b[1;31m# Tell the ConcreteFunction to clean up its graph once it goes out of\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mf:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\framework\\func_graph.py\u001b[0m in \u001b[0;36mfunc_graph_from_py_func\u001b[1;34m(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)\u001b[0m\n\u001b[0;32m    976\u001b[0m                                           converted_func)\n\u001b[0;32m    977\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 978\u001b[1;33m       \u001b[0mfunc_outputs\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mpython_func\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0mfunc_args\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mfunc_kwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    979\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    980\u001b[0m       \u001b[1;31m# invariant: `func_outputs` contains only Tensors, CompositeTensors,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mf:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\eager\\def_function.py\u001b[0m in \u001b[0;36mwrapped_fn\u001b[1;34m(*args, **kwds)\u001b[0m\n\u001b[0;32m    437\u001b[0m         \u001b[1;31m# __wrapped__ allows AutoGraph to swap in a converted function. We give\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    438\u001b[0m         \u001b[1;31m# the function a weak reference to itself to avoid a reference cycle.\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 439\u001b[1;33m         \u001b[1;32mreturn\u001b[0m \u001b[0mweak_wrapped_fn\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m__wrapped__\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    440\u001b[0m     \u001b[0mweak_wrapped_fn\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mweakref\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mref\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mwrapped_fn\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    441\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mf:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\framework\\func_graph.py\u001b[0m in \u001b[0;36mwrapper\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    966\u001b[0m           \u001b[1;32mexcept\u001b[0m \u001b[0mException\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[1;33m:\u001b[0m  \u001b[1;31m# pylint:disable=broad-except\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    967\u001b[0m             \u001b[1;32mif\u001b[0m \u001b[0mhasattr\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0me\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m\"ag_error_metadata\"\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 968\u001b[1;33m               \u001b[1;32mraise\u001b[0m \u001b[0me\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mag_error_metadata\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mto_exception\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0me\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    969\u001b[0m             \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    970\u001b[0m               \u001b[1;32mraise\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mAttributeError\u001b[0m: in converted code:\n\n    <ipython-input-16-f1680ec5eaf9>:17 train_step  *\n        per_example_losses = mirrored_strategy.run(step_fn, args=(dist_inputs,))\n\n    AttributeError: 'MirroredStrategy' object has no attribute 'run'\n"
     ]
    }
   ],
   "source": [
    "with mirrored_strategy.scope():\n",
    "  for inputs in dist_dataset:\n",
    "    print(train_step(inputs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在上面的例子中，我们遍历 dist_dataset ，为你的训练提供输入。我们还提供 tf.distribute.Strategy.make_experimental_numpy_dataset 来支持 numpy 输入。在调用[tf.distribute.Strategy.experimental_distribute_dataset](https://tensorflow.google.cn/api_docs/python/tf/distribute/Strategy#experimental_distribute_dataset) 之前，可以使用此API创建数据集。\n",
    "\n",
    "迭代数据的另一种方法是显式地使用迭代器。当你希望运行给定数量的步骤而不是遍历整个数据集时，你可能需要执行此操作。上面的迭代现在将被修改为首先创建一个迭代器，然后显式地调用它的next以获取输入数据。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "in converted code:\n\n    <ipython-input-16-f1680ec5eaf9>:17 train_step  *\n        per_example_losses = mirrored_strategy.run(step_fn, args=(dist_inputs,))\n\n    AttributeError: 'MirroredStrategy' object has no attribute 'run'\n",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-18-790610b637fa>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      2\u001b[0m   \u001b[0miterator\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0miter\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mdist_dataset\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      3\u001b[0m   \u001b[1;32mfor\u001b[0m \u001b[0m_\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m10\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 4\u001b[1;33m     \u001b[0mprint\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mtrain_step\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mnext\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0miterator\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;32mf:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\eager\\def_function.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, *args, **kwds)\u001b[0m\n\u001b[0;32m    566\u001b[0m         \u001b[0mxla_context\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mExit\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    567\u001b[0m     \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 568\u001b[1;33m       \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_call\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    569\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    570\u001b[0m     \u001b[1;32mif\u001b[0m \u001b[0mtracing_count\u001b[0m \u001b[1;33m==\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_get_tracing_count\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mf:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\eager\\def_function.py\u001b[0m in \u001b[0;36m_call\u001b[1;34m(self, *args, **kwds)\u001b[0m\n\u001b[0;32m    604\u001b[0m       \u001b[1;31m# In this case we have not created variables on the first call. So we can\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    605\u001b[0m       \u001b[1;31m# run the first trace but we should fail if variables are created.\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 606\u001b[1;33m       \u001b[0mresults\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_stateful_fn\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    607\u001b[0m       \u001b[1;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_created_variables\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    608\u001b[0m         raise ValueError(\"Creating variables on a non-first call to a function\"\n",
      "\u001b[1;32mf:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\eager\\function.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   2360\u001b[0m     \u001b[1;34m\"\"\"Calls a graph function specialized to the inputs.\"\"\"\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2361\u001b[0m     \u001b[1;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_lock\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 2362\u001b[1;33m       \u001b[0mgraph_function\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mkwargs\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_maybe_define_function\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   2363\u001b[0m     \u001b[1;32mreturn\u001b[0m \u001b[0mgraph_function\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_filtered_call\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m  \u001b[1;31m# pylint: disable=protected-access\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2364\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mf:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\eager\\function.py\u001b[0m in \u001b[0;36m_maybe_define_function\u001b[1;34m(self, args, kwargs)\u001b[0m\n\u001b[0;32m   2701\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2702\u001b[0m       \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_function_cache\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmissed\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0madd\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mcall_context_key\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 2703\u001b[1;33m       \u001b[0mgraph_function\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_create_graph_function\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   2704\u001b[0m       \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_function_cache\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mprimary\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mcache_key\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mgraph_function\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2705\u001b[0m       \u001b[1;32mreturn\u001b[0m \u001b[0mgraph_function\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mf:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\eager\\function.py\u001b[0m in \u001b[0;36m_create_graph_function\u001b[1;34m(self, args, kwargs, override_flat_arg_shapes)\u001b[0m\n\u001b[0;32m   2591\u001b[0m             \u001b[0marg_names\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0marg_names\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2592\u001b[0m             \u001b[0moverride_flat_arg_shapes\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0moverride_flat_arg_shapes\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 2593\u001b[1;33m             capture_by_value=self._capture_by_value),\n\u001b[0m\u001b[0;32m   2594\u001b[0m         \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_function_attributes\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   2595\u001b[0m         \u001b[1;31m# Tell the ConcreteFunction to clean up its graph once it goes out of\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mf:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\framework\\func_graph.py\u001b[0m in \u001b[0;36mfunc_graph_from_py_func\u001b[1;34m(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)\u001b[0m\n\u001b[0;32m    976\u001b[0m                                           converted_func)\n\u001b[0;32m    977\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 978\u001b[1;33m       \u001b[0mfunc_outputs\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mpython_func\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0mfunc_args\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mfunc_kwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    979\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    980\u001b[0m       \u001b[1;31m# invariant: `func_outputs` contains only Tensors, CompositeTensors,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mf:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\eager\\def_function.py\u001b[0m in \u001b[0;36mwrapped_fn\u001b[1;34m(*args, **kwds)\u001b[0m\n\u001b[0;32m    437\u001b[0m         \u001b[1;31m# __wrapped__ allows AutoGraph to swap in a converted function. We give\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    438\u001b[0m         \u001b[1;31m# the function a weak reference to itself to avoid a reference cycle.\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 439\u001b[1;33m         \u001b[1;32mreturn\u001b[0m \u001b[0mweak_wrapped_fn\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m__wrapped__\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    440\u001b[0m     \u001b[0mweak_wrapped_fn\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mweakref\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mref\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mwrapped_fn\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    441\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mf:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\framework\\func_graph.py\u001b[0m in \u001b[0;36mwrapper\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    966\u001b[0m           \u001b[1;32mexcept\u001b[0m \u001b[0mException\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[1;33m:\u001b[0m  \u001b[1;31m# pylint:disable=broad-except\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    967\u001b[0m             \u001b[1;32mif\u001b[0m \u001b[0mhasattr\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0me\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m\"ag_error_metadata\"\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 968\u001b[1;33m               \u001b[1;32mraise\u001b[0m \u001b[0me\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mag_error_metadata\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mto_exception\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0me\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    969\u001b[0m             \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    970\u001b[0m               \u001b[1;32mraise\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mAttributeError\u001b[0m: in converted code:\n\n    <ipython-input-16-f1680ec5eaf9>:17 train_step  *\n        per_example_losses = mirrored_strategy.run(step_fn, args=(dist_inputs,))\n\n    AttributeError: 'MirroredStrategy' object has no attribute 'run'\n"
     ]
    }
   ],
   "source": [
    "with mirrored_strategy.scope():\n",
    "  iterator = iter(dist_dataset)\n",
    "  for _ in range(10):\n",
    "    print(train_step(next(iterator)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这涵盖了使用 tf.distribute.Strategy API 分发自定义训练循环的最简单情况。tensorflow 正在改进这些 API。由于这个用例需要更多的工作来适应你的代码，tensorflow 将在未来发布一个单独的详细指南。\n",
    "\n",
    "### 现在支持什么？\n",
    "在 TF2.0 版本中，使用上面所示的 MirroredStrategy 和 TPUStrategy 支持使用自定义训练循环的训练。未来将提供 MultiWorkerMirorredStrategy 支持。\n",
    "\n",
    "| 训练 API | MirroredStrategy | TPUStrategy | MultiWorkerMirroredStrategy | CentralStorageStrategy | ParameterServerStrategy | OneDeviceStrategy |\n",
    "| - | - | - | - | - | - | - |\n",
    "| Custom Training Loop | 实验性支持 | 实验性支持 | 计划 2.0 后支持 | 计划 2.0 后支持 | 目前不支持 | 支持 |\n",
    "\n",
    "### 示例和教程\n",
    "下面是一些将分布策略用于自定义训练循环的示例：\n",
    "1. 使用镜像策略培训 MNIST 的 [教程](https://tensorflow.google.cn/tutorials/distribute/custom_training)。\n",
    "2. 使用 MirroredStrategy 的 [DenseNet](https://github.com/tensorflow/examples/blob/master/tensorflow_examples/models/densenet/distributed_train.py) 示例。\n",
    "3. [BERT](https://github.com/tensorflow/models/blob/master/official/nlp/bert/run_classifier.py) 示例使用 MirroredStrategy 和 TPUStrategy 进行训练。这个例子特别有助于理解如何从一个检查点加载，以及如何在分布式训练期间生成周期性的检查点等。\n",
    "4. 使用 MirroredStrategy 和 TPUStrategy 训练的 [NCF](https://github.com/tensorflow/models/blob/master/official/recommendation/ncf_keras_main.py) 示例，可以使用 keras_use_ctl 标志启用。\n",
    "5. 使用镜像策略训练的 [NMT](https://github.com/tensorflow/examples/blob/master/tensorflow_examples/models/nmt_with_attention/distributed_train.py) 示例。\n",
    "\n",
    "## 在 Estimator 使用 tf.distribute.Strategy（支持受限）\n",
    "[tf.estimator](https://tensorflow.google.cn/api_docs/python/tf/estimator) 是一个分布式训练 tensorflow API，最初支持异步参数服务器方法。和 Keras一样，我们将 tf.distribute.Strategy 集成到  tf.Estimator 中。如果你正在使用 Estimator 进行训练，你可以很容易地更改为分布式培训，只需对代码进行很少的更改。有了这个，估计器用户现在可以在多个 GPU 和多个工作机上进行同步分布式培训，也可以使用 TPU。然而，这种对 Estimator 的支持是有限的。\n",
    "\n",
    "带有 Estimator 的 tf.distribute.Strategy 的使用与 Keras 的情况略有不同。现在我们不使用 strategy.scope，而是将 strategy 对象传递到 Estimator 的 [RunConfig](https://tensorflow.google.cn/api_docs/python/tf/estimator/RunConfig) 中。\n",
    "\n",
    "下面是一段代码，它使用预编码估计器 LinearRegressor 和 MirroredStrategy 展示了这一点："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:There are non-GPU devices in `tf.distribute.Strategy`, not using nccl allreduce.\n",
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0',)\n",
      "INFO:tensorflow:Initializing RunConfig with distribution strategies.\n",
      "INFO:tensorflow:Not using Distribute Coordinator.\n",
      "WARNING:tensorflow:Using temporary folder as model directory: C:\\Users\\hu'nan\\AppData\\Local\\Temp\\tmpaf07vnab\n",
      "INFO:tensorflow:Using config: {'_model_dir': \"C:\\\\Users\\\\hu'nan\\\\AppData\\\\Local\\\\Temp\\\\tmpaf07vnab\", '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': allow_soft_placement: true\n",
      "graph_options {\n",
      "  rewrite_options {\n",
      "    meta_optimizer_iterations: ONE\n",
      "  }\n",
      "}\n",
      ", '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': <tensorflow.python.distribute.mirrored_strategy.MirroredStrategy object at 0x00000200C07F4148>, '_device_fn': None, '_protocol': None, '_eval_distribute': <tensorflow.python.distribute.mirrored_strategy.MirroredStrategy object at 0x00000200C07F4148>, '_experimental_distribute': None, '_experimental_max_worker_delay_secs': None, '_session_creation_timeout_secs': 7200, '_service': None, '_cluster_spec': ClusterSpec({}), '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_distribute_coordinator_mode': None}\n"
     ]
    }
   ],
   "source": [
    "mirrored_strategy = tf.distribute.MirroredStrategy()\n",
    "config = tf.estimator.RunConfig(\n",
    "    train_distribute=mirrored_strategy, eval_distribute=mirrored_strategy)\n",
    "regressor = tf.estimator.LinearRegressor(\n",
    "    feature_columns=[tf.feature_column.numeric_column('feats')],\n",
    "    optimizer='SGD',\n",
    "    config=config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们在这里使用一个预先定义的 Estimator，但是相同的代码也可以使用一个自定义的 Estimator。训练分布决定如何分配训练，评估分布决定如何分配评估。这是与 Keras 的另一个不同之处，我们在训练和评估中使用相同的策略。\n",
    "\n",
    "现在我们可以用一个输入函数来训练和评估这个 Estimator："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From f:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\ops\\resource_variable_ops.py:1635: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "If using Keras pass *_constraint arguments to layers.\n",
      "INFO:tensorflow:Reduce to /replica:0/task:0/device:CPU:0 then broadcast to ('/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "WARNING:tensorflow:From f:\\python\\pythonenv\\machine_learning\\lib\\site-packages\\tensorflow_core\\python\\feature_column\\feature_column_v2.py:518: Layer.add_variable (from tensorflow.python.keras.engine.base_layer) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Please use `layer.add_weight` method instead.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Create CheckpointSaverHook.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "INFO:tensorflow:Saving checkpoints for 0 into C:\\Users\\hu'nan\\AppData\\Local\\Temp\\tmpaf07vnab\\model.ckpt.\n",
      "INFO:tensorflow:loss = 1.0, step = 0\n",
      "INFO:tensorflow:Saving checkpoints for 10 into C:\\Users\\hu'nan\\AppData\\Local\\Temp\\tmpaf07vnab\\model.ckpt.\n",
      "INFO:tensorflow:Loss for final step: 2.877698e-13.\n",
      "INFO:tensorflow:Reduce to /replica:0/task:0/device:CPU:0 then broadcast to ('/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Reduce to /replica:0/task:0/device:CPU:0 then broadcast to ('/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /replica:0/task:0/device:CPU:0 then broadcast to ('/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /replica:0/task:0/device:CPU:0 then broadcast to ('/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /replica:0/task:0/device:CPU:0 then broadcast to ('/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /replica:0/task:0/device:CPU:0 then broadcast to ('/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /replica:0/task:0/device:CPU:0 then broadcast to ('/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /replica:0/task:0/device:CPU:0 then broadcast to ('/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /replica:0/task:0/device:CPU:0 then broadcast to ('/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Starting evaluation at 2020-04-28T20:48:16Z\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Restoring parameters from C:\\Users\\hu'nan\\AppData\\Local\\Temp\\tmpaf07vnab\\model.ckpt-10\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "INFO:tensorflow:Evaluation [1/10]\n",
      "INFO:tensorflow:Evaluation [2/10]\n",
      "INFO:tensorflow:Evaluation [3/10]\n",
      "INFO:tensorflow:Evaluation [4/10]\n",
      "INFO:tensorflow:Evaluation [5/10]\n",
      "INFO:tensorflow:Evaluation [6/10]\n",
      "INFO:tensorflow:Evaluation [7/10]\n",
      "INFO:tensorflow:Evaluation [8/10]\n",
      "INFO:tensorflow:Evaluation [9/10]\n",
      "INFO:tensorflow:Evaluation [10/10]\n",
      "INFO:tensorflow:Inference Time : 0.49169s\n",
      "INFO:tensorflow:Finished evaluation at 2020-04-28-20:48:16\n",
      "INFO:tensorflow:Saving dict for global step 10: average_loss = 1.4210855e-14, global_step = 10, label/mean = 1.0, loss = 1.4210855e-14, prediction/mean = 0.99999994\n",
      "INFO:tensorflow:Saving 'checkpoint_path' summary for global step 10: C:\\Users\\hu'nan\\AppData\\Local\\Temp\\tmpaf07vnab\\model.ckpt-10\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'average_loss': 1.4210855e-14,\n",
       " 'label/mean': 1.0,\n",
       " 'loss': 1.4210855e-14,\n",
       " 'prediction/mean': 0.99999994,\n",
       " 'global_step': 10}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def input_fn():\n",
    "  dataset = tf.data.Dataset.from_tensors(({\"feats\":[1.]}, [1.]))\n",
    "  return dataset.repeat(1000).batch(10)\n",
    "regressor.train(input_fn=input_fn, steps=10)\n",
    "regressor.evaluate(input_fn=input_fn, steps=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这里要强调的是 Estimator 和 Keras 的另一个区别是输入处理。在 Keras 中，我们提到了数据集的每一批都是在多个副本之间自动分割的。然而，在 Estimator 中，我们不会自动分割批处理，也不会自动在不同的工作机之间切分数据。你可以完全控制如何将数据分布到工作机和设备上，并且必须提供一个输入来指定如何分布数据。\n",
    "\n",
    "每个工作机调用一次输入，从而为每个工作机提供一个数据集。然后，该数据集的一个批被馈送到该工作机上的一个副本，从而在 1 个工作机上为 N 个副本消耗 N 批。换句话说，input_fn 返回的数据集应该提供每个副本的批大小。一个步骤的全局批量大小可以根据同步中的REPLICA_batch_size * strategy.num_replicas_in_sync。\n",
    "\n",
    "在进行多工作机训练时，你应该将数据分散到多个工作机中，或者在每个员工中随机抽取一个种子。你可以在 [Multi-worker Training with Estimator](https://tensorflow.google.cn/tutorials/distribute/multi_worker_with_estimator) 中看到一个如何做到这一点的示例。\n",
    "\n",
    "我们给出了一个使用带 Estimator 的镜像策略的例子。你还可以将 TPUStrategy 与 Estimator 一起使用，方法完全相同："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'tpu_strategy' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-21-dd90b675713b>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      1\u001b[0m config = tf.estimator.RunConfig(\n\u001b[1;32m----> 2\u001b[1;33m     train_distribute=tpu_strategy, eval_distribute=tpu_strategy)\n\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m: name 'tpu_strategy' is not defined"
     ]
    }
   ],
   "source": [
    "config = tf.estimator.RunConfig(\n",
    "    train_distribute=tpu_strategy, eval_distribute=tpu_strategy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "同样，你也可以使用多工作机和参数服务器策略。代码保持不变，但你需要使用  [tf.estimator.train_and_evaluate](https://tensorflow.google.cn/api_docs/python/tf/estimator/train_and_evaluate)，并为集群中运行的每个二进制文件设置 tf_CONFIG 环境变量。\n",
    "\n",
    "### 现在支持什么？\n",
    "在 TF2.0 版本中，除了 TPUStrategy 之外，对使用所有策略的估计器进行培训的支持是有限的。基本的训练和评估是有效的，但是一些先进的功能，如脚手架还没有使用。在这个集成中也可能有一些 bug。目前，tensorflow 不打算积极改善这种支持，而是专注于 Keras 和定制训练循环支持。如果可能的话，你应该更愿意将 tf.distribute 与这些 API 一起使用。\n",
    "\n",
    "| 训练 API | MirroredStrategy | TPUStrategy | MultiWorkerMirroredStrategy | CentralStorageStrategy | ParameterServerStrategy | OneDeviceStrategy |\n",
    "| - | - | - | - | - | - | - |\n",
    "| Estimator API | 支持受限 | 不支持 | 支持受限 |支持受限 | 支持受限 | 支持受限 |\n",
    "\n",
    "### 示例和教程\n",
    "下面是一些示例，显示了使用 Estimator 的各种策略的端到端用法：\n",
    "- [Multi-worker Training with Estimator ](https://tensorflow.google.cn/tutorials/distribute/multi_worker_with_estimator) 用 MultiWorkerMirroredStrategy 训练多工作机的 MNIST。\n",
    "- 使用 Kubernetes 模板在 tensorflow/生态系统 中进行多工作机培训的端到端示例。本例从 Keras 模型开始，并使用 [tf.Keras.Estimator.model_to_Estimator](https://tensorflow.google.cn/api_docs/python/tf/keras/estimator/model_to_estimator) API将其转换为 Estimator。\n",
    "- 正式的 [ResNet50](https://github.com/tensorflow/models/blob/master/official/vision/image_classification/resnet_imagenet_main.py) 模型，可以使用 MirroredStrategy 或 MultiWorkerMirroredStrategy 进行训练。\n",
    "\n",
    "## 其它主题\n",
    "### 设置 TF_CONFIG 环境变量\n",
    "如前所述，对于多工作机训练，需要为集群中运行的每个二进制文件设置 TF_CONFIG 环境变量。TF_CONFIG 环境变量是一个 JSON 字符串，指定哪些任务构成集群、它们的地址以及每个任务在集群中的角色。我们在 [tensorflow/economic repo](https://github.com/tensorflow/ecosystem) 中提供了一个 Kubernetes 模板，它为你的培训练任务设置 TF_CONFIG。\n",
    "\n",
    "TF_CONFIG 的一个例子是：\n",
    "```json\n",
    "os.environ[\"TF_CONFIG\"] = json.dumps({\n",
    "    \"cluster\": {\n",
    "        \"worker\": [\"host1:port\", \"host2:port\", \"host3:port\"],\n",
    "        \"ps\": [\"host4:port\", \"host5:port\"]\n",
    "    },\n",
    "   \"task\": {\"type\": \"worker\", \"index\": 1}\n",
    "})\n",
    "```\n",
    "此 TF_CONFIG 指定集群中有三个工作机和两个 ps 任务及其主机和端口。\"task\" 部分指定集群中当前任务的角色 worker1（第二个 worker）。群集中的有效角色是 \"主管\"、\"工人\"、\"ps\" 和 \"评估者\"。除了使用tf.distribute.experiential.ParameterServerStrategy 时，不应该有 \"ps\" 任务。\n",
    "\n",
    "## 下一步\n",
    "tf.distribute.Strategy 正在积极制定中。tensorflow 欢迎你使用进行尝试并提供和反馈于 [GitHub issue](https://github.com/tensorflow/tensorflow/issues/new)。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6rc1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
